package com.fengwk.cv4j.utils;

import java.util.Arrays;

import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

import com.fengwk.support.exception.BaseException;
import com.fengwk.support.function.Return;
import com.fengwk.support.util.ArrayUtils;

/**
 * Nd4j相关工具
 * 
 * @author fengwk
 *
 */
public class Nd4jUtils {

	/**
	 * 回退式的广播</br>
	 * 在INDArray的广播规则中并不能从低维数向高维数扩张</br>
	 * 这个方法支持例如2*3向3*(2*3)这种形式的广播</br>
	 * 这样的广播实质上是将低纬度数组堆叠成一个高纬度数组的广播
	 * 
	 * @param src
	 * @return
	 */
	public static INDArray backBroadcast(INDArray src, int...shape) {
		int[] shape0 = src.shape();
		if (shape0.length > shape.length)
			throw new BaseException("shape0.length > shape.length");
		if (shape0.length == shape.length)
			return src.broadcast(shape);
		// 拓展维度
		for (int i = 0; i < shape.length - shape0.length; i ++)
			src = Nd4j.expandDims(src, 0);
		return src.broadcast(shape);
	}
	
	/**
	 * 最大值,使维度相对保持</br>
	 * {@link #keepDimension(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray max(INDArray arr, int...dimension) {
		return keepDimension((a, d) -> a.max(d), arr, dimension);
	}
	
	/**
	 * 最小值,使维度相对保持</br>
	 * {@link #keepDimension(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray min(INDArray arr, int...dimension) {
		return keepDimension((a, d) -> a.min(d), arr, dimension);
	}
	
	/**
	 * 求和,使维度相对保持</br>
	 * {@link #keepDimension(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray sum(INDArray arr, int...dimension) {
		return keepDimension((a, d) -> a.sum(d), arr, dimension);
	}
	
	/**
	 * 求均值,使维度相对保持</br>
	 * {@link #keepDimension(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray mean(INDArray arr, int...dimension) {
		return keepDimension((a, d) -> a.mean(d), arr, dimension);
	}
	
	/**
     * 保持相对维度</br>
     * 在INDArray的维度相关运算中:大于2个维度的运算会产生高维度丢失,而在小于大于2个维度时维度却不会丢失</br>
     * 这个方法的作用在于保持INDArray的维度运算的规则,并使其不出现一些副作用,例如:</br>
     * (2,2,3)的张量进行维度(0,2)维度运算后会变成一个(1,2)的矩阵,而大多数情况下,我们需要一个(2,1)的矩阵</br>
     * 这个方法就是为了确保在维度运算时能呈现出如上我们想要的维度
     * 
     * @param arr
     * @param dimension
     * @return
     */
    private static INDArray keepDimension(Return.Params2<INDArray, INDArray, int[]> op, INDArray arr, int...dimension) {
    	if (ArrayUtils.isBlank(dimension))
    		throw new BaseException("ArrayUtils.isBlank(dimension)");
    	
    	if (arr.rank() <= 2)
    		return op.apply(arr, dimension);
    	// 维度从小到大逐个操作,使维度保持
    	Arrays.sort(dimension);
    	int[] d0 = new int[dimension.length < 2 ? dimension.length : 2];
    	int offset = 0;
    	for (int i = 0; i < dimension.length; i ++) {
    		if (i < dimension.length - 2) 
    			arr = op.apply(arr, new int[] {dimension[i] - (offset ++)});
    		else 
    			d0[i - dimension.length + d0.length] = dimension[i] - offset;
    	}
    	return op.apply(arr, d0);
    }
    
    /**
	 * 最大值,使维度保持</br>
	 * {@link #keepDimension0(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray max0(INDArray arr, int...dimension) {
		return keepDimension0((a, d) -> a.max(d), arr, dimension);
	}
	
	/**
	 * 最小值,使维度保持</br>
	 * {@link #keepDimension0(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray min0(INDArray arr, int...dimension) {
		return keepDimension0((a, d) -> a.min(d), arr, dimension);
	}
	
	/**
	 * 求和,使维度保持</br>
	 * {@link #keepDimension0(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray sum0(INDArray arr, int...dimension) {
		return keepDimension0((a, d) -> a.sum(d), arr, dimension);
	}
	
	/**
	 * 求均值,使维度保持</br>
	 * {@link #keepDimension0(Return.Params2, INDArray, int...)}
	 * 
	 * @param arr
	 * @param dimension
	 * @return
	 */
	public static INDArray mean0(INDArray arr, int...dimension) {
		return keepDimension0((a, d) -> a.mean(d), arr, dimension);
	}
    
    /**
     * 保持维度</br>
     * 与{@link #keepDimension(Return.Params2, INDArray, int...)}不同的是,这个方法确实地保持了维度的准确性,不会造成维度的丢失</br>
     * 在运算中保持维度准确通常是为了精准的广播
     * 
     * @param op
     * @param arr
     * @param dimension
     * @return
     */
    private static INDArray keepDimension0(Return.Params2<INDArray, INDArray, int[]> op, INDArray arr, int...dimension) {
    	if (ArrayUtils.isBlank(dimension))
    		throw new BaseException("ArrayUtils.isBlank(dimension)");
    	
    	if (arr.rank() <= 2)
    		return op.apply(arr, dimension);
    	Arrays.sort(dimension);
    	for (int i = 0; i < dimension.length; i ++) {
			arr = op.apply(arr, new int[] {dimension[i]});
			arr = Nd4j.expandDims(arr, dimension[i]);
    	}
    	return arr;
    }
    
}
